{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 使用 SWIFT 进行推理\n",
    "\n",
    "SWIFT 是 ModelScope 开源的一款框架，支持大模型的训练、推理、评测和部署。通过 SWIFT，可以轻松实现从模型训练到应用的完整链路。\n",
    "\n",
    "本教程将详细介绍如何使用 SWIFT 进行推理，包括安装步骤和推理示例。我们将以 Yi-1.5-6B-Chat 为例进行演示。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🚀 使用 Colab 运行\n",
    "\n",
    "我们提供了一键运行的 [Colab 脚本](https://colab.research.google.com/drive/1R0s7cDNWTNCWjod_z-jVpxiFc-R3_7kc?usp=drive_link)，你可以直接在 Colab 中运行这个教程。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 安装\n",
    "\n",
    "首先，我们需要安装相关的依赖。\n",
    "\n",
    "（可选）可以选择设置 pip 全局镜像以加快下载速度：\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "!pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "安装 ms-swift：\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "!pip install 'ms-swift[llm]' -U"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 开始推理\n",
    "\n",
    "在开始推理之前，需要注意你的电脑显存和内存。如果显存不足，可能会导致报错。\n",
    "\n",
    "| 模型           | 显存使用 | 硬盘占用 |\n",
    "| -------------- | -------- | -------- |\n",
    "| Yi-1.5-6B-Chat | 11.5G    | 14.7G    |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们首先设置环境变量："
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后加载模型和分词器："
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from swift.llm import (\n",
    "    get_model_tokenizer, get_template, inference, ModelType, get_default_template_type,\n",
    ")\n",
    "from swift.utils import seed_everything\n",
    "\n",
    "# 选择模型类型，这里使用 Yi-1.5-6B-Chat\n",
    "model_type = ModelType.yi_1_5_6b_chat\n",
    "template_type = get_default_template_type(model_type)\n",
    "print(f'template_type: {template_type}')  # 模板类型\n",
    "\n",
    "# 加载模型和分词器\n",
    "kwargs = {}\n",
    "model, tokenizer = get_model_tokenizer(model_type, model_kwargs={'device_map': 'auto'}, **kwargs)\n",
    "\n",
    "# 设置生成配置\n",
    "model.generation_config.max_new_tokens = 128\n",
    "\n",
    "# 获取模板\n",
    "template = get_template(template_type, tokenizer)\n",
    "\n",
    "# 设置随机种子\n",
    "seed_everything(42)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来，我们进行推理："
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# 准备输入查询\n",
    "query = '你好！'\n",
    "\n",
    "# 使用模板进行推理\n",
    "response, history = inference(model, template, query)\n",
    "\n",
    "# 打印查询和响应\n",
    "print(f'query: {query}')\n",
    "print(f'response: {response}')"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "以上代码将会输出类似如下的结果：\n",
    "\n",
    "```\n",
    "query: 你好！\n",
    "response: 你好！很高兴见到你。有什么我可以帮忙的吗？\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "到这里，你已经学会了如何使用 SWIFT 进行 Yi 系列模型的推理。如果在使用过程中遇到任何问题，可以参考 [SWIFT 的官方文档](https://www.modelscope.cn/models/01-ai/Yi-1.5-6B-Chat) 获取更多帮助。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
